import torch
import numpy as np

from utils.Logger import get_logger
from utils.Radar_draw import plot_single_heatmap
from utils.nameUtils import generate_future_npy_filenames
from utils.rainDraw import draw_cnn_png
from utils.history import plot_rain_with_map
from utils.nameUtils import getOneHourLater, getCurrentTimeName
import yaml
import os
import shutil
import threading

logger = get_logger()  # 配置日志

# 配置类
with open('config.yaml', 'r') as file:
    config = yaml.safe_load(file)


# 可视化保存雷达预测的dbz和图片  保存结果都是模型直接推导出来的结果  基本在-1到1之间
def save_image(output: torch.tensor, config: dict, deadLine: str):
    '''

    :param config:
    :param deadLine: 当前时间戳
    :param output:(1,20,1,400,400)
    :return:
    '''
    # 预测之后的结果控制在0-1之间
    data = output[0, :, 0, ...].cpu().numpy()
    data = np.clip(data, 0, 1)
    # 未来20帧的图片保存路径
    name_list = generate_future_npy_filenames(config['image']['dbz']['save_path'], deadLine, "png", 20)
    npy_list = generate_future_npy_filenames(config['radar']['save_path'], deadLine, "npy", 20)
    backupNameList = generate_future_npy_filenames(config['radar']['dbz_backup'], deadLine, "npy", 20)
    for i, name in enumerate(name_list):
        np.save(npy_list[i], data[i])
        # 备份DBZ 只需要备份最大值大于25的文件
        if data[i].max() * 70 > 25:
            logger.info("备份dbz文件")
            np.save(backupNameList[i], data[i])
        plot_single_heatmap(data[i] * 70, name, "png", transparent=True)


# 保存降雨估计模型的npy文件
def save_rain_npy(inputFrame, filename, config):
    '''

    :param inputFrame:待保存的文件  numpy格式
    :param filename: 存放的文件名
    :return:
    @param config:
    '''
    # 限制inputFrame的值都是大于0
    np.maximum(inputFrame, 0)
    # 这里是直接传入的202301010303.png格式 所以需要去掉后缀
    np.save(filename, inputFrame)


def turn_dbz_to_zr(input_dbz, npy_path, png_path, backupPath):
    '''
    Z-R关系转化  直接将dbz转化为降雨量
    :param input_dbz: 0-1之间的dbz
    :param filename: 保存的文件路径
    :return:
    @param backupPath: 备份保存目录
    '''
    # 后面转化之前要先乘以70
    input_dbz = input_dbz * 70
    input_dbz = np.where(input_dbz < 15, 0, (10 ** (input_dbz / 10) / 10 ** 2.12359) ** (1 / 1.23418))
    # 保存降雨可视化结果
    draw_cnn_png(input_dbz, png_path)
    # 保存降雨结果
    np.save(npy_path, input_dbz)
    # 备份降水值 只保存大于15的降雨
    if input_dbz.max() > 15:
        logger.info("备份降雨文件")
        np.save(backupPath, input_dbz)


# 保存历史降水记录
def save_history(input_dbz, lastName):
    '''

    :param input_dbz: (20,1,400,400)
    :param lastName:
    :return:
    '''
    # 转化为ZR
    input_dbz = input_dbz * 70
    input_dbz = np.where(input_dbz < 15, 0, (10 ** (input_dbz / 10) / 10 ** 2.12359) ** (1 / 1.23418))
    oneHourData = np.sum(input_dbz[:10], axis=0)
    path = os.path.join(config["historyPath"], lastName)
    if not os.path.exists(path):
        os.makedirs(path, exist_ok=True)
    twoHourData = np.sum(input_dbz, axis=0)
    path2 = os.path.join(config["historyPath"], lastName)
    if not os.path.exists(path2):
        os.makedirs(path2, exist_ok=True)
    logger.info("保存未来" + lastName + "预测1小时的数据")
    # threading.Thread(target=plot_rain_with_map, args=(oneHourData, getCurrentTimeName(lastName) + "未来一个小时的降雨预估", path + "/oneHour.png", "png")).start()
    plot_rain_with_map(oneHourData, title=getCurrentTimeName(lastName) + "未来一个小时的降雨预估", fname=path + "/oneHour.png", format="png")
    logger.info("保存未来" + lastName + "预测2小时的数据")
    # threading.Thread(target=plot_rain_with_map, args=(twoHourData, getCurrentTimeName(lastName) + "未来两个小时的降雨预估", path2 + "/twoHour.png", "png")).start()
    plot_rain_with_map(twoHourData, title=getCurrentTimeName(lastName) + "未来两个小时的降雨预估", fname=path2 + "/twoHour.png", format="png")
    # 把当前统计的降水量保存到默认显示位置
    shutil.copy2(path + "/oneHour.png", config["historyPath"] + "/oneHour.png")
    shutil.copy2(path2 + "/twoHour.png", config["historyPath"] + "/twoHour.png")
